from langchain.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langchain_neo4j import Neo4jGraph, GraphCypherQAChain
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, Type, Optional
from dotenv import load_dotenv
from model.my_chat_model import ChatModel
import os


class Neo4jInput(BaseModel):
    question: str = Field(..., description="问题")


class Neo4jTool(BaseTool):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __init__(self, **kwargs: Any):
        super().__init__(
            name="get_neo4j_tool",
            description="主要用于查询知识图谱的数据，例如人物，公司的关系数据，必须输入的参数是问题",
            **kwargs
        )

    args_schema: Type[BaseModel] = Neo4jInput

    def _run(self, question: str):
        load_dotenv()
        chat = ChatModel()
        llm = chat.get_line_model()
        graph = Neo4jGraph(url=os.getenv("NEO4J_URI"), username=os.getenv("NEO4J_USERNAME"),
                           password=os.getenv("NEO4J_PASSWORD"))

        # 核心修正：Prompt 明确聚合函数使用规则，新增正确的股东查询示例
        prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                """你是专业的 Neo4j Cypher 查询生成器，仅输出纯 Cypher 语句（无注释、无多余文本、无解释），严格遵循以下规则：

                一、数据库结构（不变）
                1. 节点标签：企业(:企业)、十大股东(:十大股东)、高管(:高管)、最终受益人(:最终受益人)；
                2. 关系：十大股东-[:持股]->企业，高管-[:任职于]->企业；
                3. 核心属性：企业.企业名称、企业.报告期、十大股东.持股比例、十大股东.股东排名。

                二、Cypher 生成规则（新增聚合函数语法限制）
                1. 模糊匹配企业：用 `c.企业名称 CONTAINS "关键词"`（如“茅台”），禁止精确匹配除非用户输入全称；
                2. 聚合函数规则（重点！）：
                   - 若需用 MAX(c.报告期) 获取最新报告期，**必须先用 WITH 子句计算**，禁止直接在 WHERE 中使用；
                   - 正确流程：先匹配企业→用 WITH 算最新报告期→再关联其他节点（如股东）→最后过滤报告期；
                3. 多节点关联（如查股东）：必须先确定企业和最新报告期，再关联股东-持股关系，禁止直接关联后用聚合函数；
                4. 返回核心属性：查股东仅返回股东名称、持股比例、股东排名、企业名称、报告期。

                三、典型需求示例（重点补充股东查询的正确语法）
                1. 需求：“茅台的十大股东信息（最新报告期）”
                   正确 Cypher：
                   MATCH (c:企业) 
                   WHERE c.企业名称 CONTAINS "茅台"  // 第一步：模糊匹配企业
                   WITH c, MAX(c.报告期) AS latest_period  // 第二步：用 WITH 算最新报告期
                   MATCH (s:十大股东)-[:持股]->(c:企业)  // 第三步：关联股东-企业关系
                   WHERE c.企业名称 CONTAINS "茅台" AND c.报告期 = latest_period  // 第四步：过滤最新报告期
                   RETURN s.股东名称, s.持股比例, s.股东排名, c.企业名称, c.报告期
                   ORDER BY s.股东排名 ASC

                2. 需求：“茅台2025-06-30的十大股东”（指定报告期）
                   正确 Cypher：
                   MATCH (s:十大股东)-[:持股]->(c:企业)
                   WHERE c.企业名称 CONTAINS "茅台" AND c.报告期 = "2025-06-30"
                   RETURN s.股东名称, s.持股比例, s.股东排名, c.企业名称

                3. 需求：“茅台最新报告期的净利润”
                   正确 Cypher：
                   MATCH (c:企业)
                   WHERE c.企业名称 CONTAINS "茅台"
                   WITH c, MAX(c.报告期) AS latest_period
                   MATCH (c:企业)
                   WHERE c.企业名称 CONTAINS "茅台" AND c.报告期 = latest_period
                   RETURN c.企业名称, c.报告期, c.净利润

                四、禁止规则
                - 禁止在 WHERE 子句中直接使用 MAX()/MIN() 等聚合函数；
                - 禁止关联节点后不过滤企业名称（如仅 MATCH (s)-[:持股]->(c) 不写 WHERE c.企业名称...）；
                - 禁止返回无关属性（如查股东返回十大股东.出生年月）。
                """),
            ("human", "{question}"),
        ])

        chain = GraphCypherQAChain.from_llm(
            llm=llm,
            graph=graph,
            allow_dangerous_requests=True,
            cypher_prompt=prompt,
            verbose=True,
            top_k=10,
            validate_cypher=True
        )

        # 捕获 Cypher 语法错误，返回友好提示
        try:
            rs = chain.invoke({"query": question})
            return rs["result"]
        except Exception as e:
            error_msg = str(e)
            if "SyntaxError" in error_msg and "aggregating function" in error_msg:
                return f"Cypher 语法错误：聚合函数（如 MAX）不能直接用在 WHERE 中，请参考示例用 WITH 子句计算后再过滤"
            else:
                return f"查询失败：{error_msg[:80]}"